package mikey

import (
	"testing"

	"github.com/stretchr/testify/require"
)

var cases = []struct {
	name string
	enc  []byte
	dec  Message
}{
	{
		"a",
		[]byte{
			0x01, 0x00, 0x05, 0x00, 0xe6, 0x9d, 0x51, 0xf8,
			0x01, 0x00, 0x00, 0x30, 0x68, 0x57, 0x60, 0x00,
			0x00, 0x00, 0x00, 0x0b, 0x00, 0xeb, 0xfe, 0x6f,
			0x2d, 0xb1, 0xc1, 0x3f, 0xd0, 0x0a, 0x10, 0xc2,
			0xdd, 0xe4, 0x43, 0xa8, 0x49, 0x30, 0xa5, 0x75,
			0x7a, 0x7e, 0xd9, 0xc3, 0xa4, 0x17, 0xfb, 0x01,
			0x00, 0x00, 0x00, 0x15, 0x00, 0x01, 0x01, 0x01,
			0x01, 0x10, 0x02, 0x01, 0x01, 0x03, 0x01, 0x0a,
			0x07, 0x01, 0x01, 0x08, 0x01, 0x01, 0x0a, 0x01,
			0x01, 0x00, 0x00, 0x00, 0x22, 0x00, 0x20, 0x00,
			0x1e, 0x90, 0x91, 0x78, 0x3d, 0xfc, 0xe8, 0xdd,
			0xcd, 0x44, 0x3a, 0x53, 0x50, 0x8b, 0x64, 0x50,
			0x9f, 0x35, 0xbd, 0x8a, 0x86, 0xbc, 0x4d, 0x8b,
			0x76, 0x37, 0xa5, 0x02, 0x49, 0x3d, 0xaf, 0x00,
		},
		Message{
			Header: Header{
				Version: 1,
				CSBID:   3869069816,
				CSIDMapInfo: []SRTPIDEntry{
					{
						PolicyNo: 0,
						SSRC:     812144480,
						ROC:      0,
					},
				},
			},
			Payloads: []Payload{
				&PayloadT{
					TSType:  0,
					TSValue: 17005151485044015056,
				},
				&PayloadRAND{
					Data: []byte{
						0xc2, 0xdd, 0xe4, 0x43, 0xa8, 0x49, 0x30, 0xa5,
						0x75, 0x7a, 0x7e, 0xd9, 0xc3, 0xa4, 0x17, 0xfb,
					},
				},
				&PayloadSP{
					PolicyParams: []PayloadSPPolicyParam{
						{
							Type:  PayloadSPPolicyParamTypeEncrAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionEncrKeyLen,
							Value: []byte{0x10},
						},
						{
							Type:  PayloadSPPolicyParamTypeAuthAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionAuthKeyLen,
							Value: []byte{0x0a},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTCPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPAuthOffOn,
							Value: []byte{1},
						},
					},
				},
				&PayloadKEMAC{
					SubPayloads: []*SubPayloadKeyData{
						{
							Type: 2,
							KeyData: []byte{
								0x90, 0x91, 0x78, 0x3d, 0xfc, 0xe8, 0xdd, 0xcd,
								0x44, 0x3a, 0x53, 0x50, 0x8b, 0x64, 0x50, 0x9f,
								0x35, 0xbd, 0x8a, 0x86, 0xbc, 0x4d, 0x8b, 0x76,
								0x37, 0xa5, 0x02, 0x49, 0x3d, 0xaf,
							},
						},
					},
				},
			},
		},
	},
	{
		"b",
		[]byte{
			0x01, 0x00, 0x05, 0x00, 0xfe, 0xaf, 0x97, 0x52,
			0x01, 0x00, 0x00, 0xcc, 0x83, 0x62, 0x37, 0x00,
			0x00, 0x00, 0x00, 0x0b, 0x00, 0xeb, 0xfe, 0xf6,
			0x6b, 0xa2, 0x8c, 0x9b, 0x84, 0x0a, 0x10, 0x27,
			0x6e, 0x94, 0x18, 0x0e, 0x88, 0x75, 0xc2, 0xea,
			0xad, 0x31, 0xd8, 0x2f, 0x86, 0x46, 0x20, 0x01,
			0x00, 0x00, 0x00, 0x15, 0x00, 0x01, 0x01, 0x01,
			0x01, 0x10, 0x02, 0x01, 0x01, 0x03, 0x01, 0x0a,
			0x07, 0x01, 0x01, 0x08, 0x01, 0x01, 0x0a, 0x01,
			0x01, 0x00, 0x00, 0x00, 0x22, 0x00, 0x20, 0x00,
			0x1e, 0x99, 0x1b, 0x0f, 0x14, 0x8f, 0x09, 0x4b,
			0x4e, 0x5b, 0x8b, 0x30, 0x53, 0xcd, 0x62, 0x76,
			0x87, 0x7f, 0xcc, 0xed, 0x18, 0x66, 0xf1, 0x41,
			0x77, 0x2a, 0xdd, 0xdd, 0xe7, 0x06, 0x4b, 0x00,
		},
		Message{
			Header: Header{
				Version: 1,
				CSBID:   4272920402,
				CSIDMapInfo: []SRTPIDEntry{
					{
						PolicyNo: 0,
						SSRC:     3431162423,
						ROC:      0,
					},
				},
			},
			Payloads: []Payload{ //nolint:dupl
				&PayloadT{
					TSValue: 17005300185146628996,
				},
				&PayloadRAND{
					Data: []byte{
						0x27, 0x6e, 0x94, 0x18, 0x0e, 0x88, 0x75, 0xc2,
						0xea, 0xad, 0x31, 0xd8, 0x2f, 0x86, 0x46, 0x20,
					},
				},
				&PayloadSP{
					PolicyParams: []PayloadSPPolicyParam{
						{
							Type:  PayloadSPPolicyParamTypeEncrAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionEncrKeyLen,
							Value: []byte{0x10},
						},
						{
							Type:  PayloadSPPolicyParamTypeAuthAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionAuthKeyLen,
							Value: []byte{0x0a},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTCPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPAuthOffOn,
							Value: []byte{1},
						},
					},
				},
				&PayloadKEMAC{
					SubPayloads: []*SubPayloadKeyData{
						{
							Type: 2,
							KeyData: []byte{
								0x99, 0x1b, 0x0f, 0x14, 0x8f, 0x09, 0x4b, 0x4e,
								0x5b, 0x8b, 0x30, 0x53, 0xcd, 0x62, 0x76, 0x87,
								0x7f, 0xcc, 0xed, 0x18, 0x66, 0xf1, 0x41, 0x77,
								0x2a, 0xdd, 0xdd, 0xe7, 0x06, 0x4b,
							},
						},
					},
				},
			},
		},
	},
	{
		"c",
		[]byte{
			0x01, 0x00, 0x05, 0x00, 0x7d, 0xe1, 0x27, 0xa6,
			0x02, 0x00, 0x00, 0xcc, 0x83, 0x62, 0x37, 0x00,
			0x00, 0x00, 0x00, 0x00, 0xb5, 0xcc, 0x3b, 0xf2,
			0x00, 0x00, 0x00, 0x00, 0x0b, 0x00, 0xeb, 0xfe,
			0xf6, 0x6b, 0xa2, 0xb1, 0xf6, 0x87, 0x0a, 0x10,
			0x61, 0xbb, 0x19, 0x94, 0x32, 0x53, 0x03, 0x56,
			0xa2, 0xd1, 0x88, 0x07, 0x15, 0x23, 0x75, 0x95,
			0x01, 0x00, 0x00, 0x00, 0x15, 0x00, 0x01, 0x01,
			0x01, 0x01, 0x10, 0x02, 0x01, 0x01, 0x03, 0x01,
			0x0a, 0x07, 0x01, 0x01, 0x08, 0x01, 0x01, 0x0a,
			0x01, 0x01, 0x00, 0x00, 0x00, 0x22, 0x00, 0x20,
			0x00, 0x1e, 0x99, 0x1b, 0x0f, 0x14, 0x8f, 0x09,
			0x4b, 0x4e, 0x5b, 0x8b, 0x30, 0x53, 0xcd, 0x62,
			0x76, 0x87, 0x7f, 0xcc, 0xed, 0x18, 0x66, 0xf1,
			0x41, 0x77, 0x2a, 0xdd, 0xdd, 0xe7, 0x06, 0x4b,
			0x00,
		},
		Message{
			Header: Header{
				Version: 1,
				CSBID:   2111907750,
				CSIDMapInfo: []SRTPIDEntry{
					{
						PolicyNo: 0,
						SSRC:     3431162423,
						ROC:      0,
					},
					{
						PolicyNo: 0,
						SSRC:     3050060786,
						ROC:      0,
					},
				},
			},
			Payloads: []Payload{ //nolint:dupl
				&PayloadT{
					TSValue: 17005300185149077127,
				},
				&PayloadRAND{
					Data: []byte{
						0x61, 0xbb, 0x19, 0x94, 0x32, 0x53, 0x03, 0x56,
						0xa2, 0xd1, 0x88, 0x07, 0x15, 0x23, 0x75, 0x95,
					},
				},
				&PayloadSP{
					PolicyParams: []PayloadSPPolicyParam{
						{
							Type:  PayloadSPPolicyParamTypeEncrAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionEncrKeyLen,
							Value: []byte{0x10},
						},
						{
							Type:  PayloadSPPolicyParamTypeAuthAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionAuthKeyLen,
							Value: []byte{0x0a},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTCPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPAuthOffOn,
							Value: []byte{1},
						},
					},
				},
				&PayloadKEMAC{
					SubPayloads: []*SubPayloadKeyData{
						{
							Type: 2,
							KeyData: []byte{
								0x99, 0x1b, 0x0f, 0x14, 0x8f, 0x09, 0x4b, 0x4e,
								0x5b, 0x8b, 0x30, 0x53, 0xcd, 0x62, 0x76, 0x87,
								0x7f, 0xcc, 0xed, 0x18, 0x66, 0xf1, 0x41, 0x77,
								0x2a, 0xdd, 0xdd, 0xe7, 0x06, 0x4b,
							},
						},
					},
				},
			},
		},
	},
	{
		"Extra padding at the end",
		[]byte{
			0x01, 0x00, 0x05, 0x00, 0xb1, 0x40, 0x56, 0x57, 0x01, 0x00,
			0x00, 0x72, 0x6e, 0xf4, 0xdb, 0x00, 0x00, 0x00, 0x00, 0x0b,
			0x00, 0xec, 0xd1, 0x50, 0x81, 0xbe, 0xdc, 0xe3, 0x97, 0x0a,
			0x10, 0xdd, 0x72, 0x24, 0x83, 0x95, 0xc2, 0xcc, 0x1a, 0x6d,
			0xe6, 0x0e, 0x42, 0x2a, 0x7f, 0x57, 0x09, 0x01, 0x00, 0x00,
			0x00, 0x18, 0x00, 0x01, 0x01, 0x01, 0x01, 0x10, 0x02, 0x01,
			0x01, 0x03, 0x01, 0x14, 0x07, 0x01, 0x01, 0x08, 0x01, 0x01,
			0x0a, 0x01, 0x01, 0x0b, 0x01, 0x0a, 0x00, 0x00, 0x00, 0x22,
			0x00, 0x20, 0x00, 0x1e, 0x5d, 0xb1, 0x8d, 0x95, 0x6f, 0x69,
			0x67, 0xcc, 0x0d, 0x73, 0xf8, 0xb4, 0xe7, 0x76, 0xa4, 0x8a,
			0x08, 0xc4, 0x22, 0x0e, 0x52, 0x1d, 0xca, 0xa3, 0xb8, 0x00,
			0xf5, 0x84, 0xfc, 0x25, 0x00, 0x00,
		},
		Message{
			Header: Header{
				Version: 1,
				CSBID:   2973783639,
				CSIDMapInfo: []SRTPIDEntry{
					{
						PolicyNo: 0,
						SSRC:     1919874267,
						ROC:      0,
					},
				},
			},
			Payloads: []Payload{
				&PayloadT{
					TSValue: 17064508981266670487,
				},
				&PayloadRAND{
					Data: []byte{
						0xdd, 0x72, 0x24, 0x83, 0x95, 0xc2, 0xcc, 0x1a, 0x6d, 0xe6,
						0x0e, 0x42, 0x2a, 0x7f, 0x57, 0x09,
					},
				},
				&PayloadSP{
					PolicyParams: []PayloadSPPolicyParam{
						{
							Type:  PayloadSPPolicyParamTypeEncrAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionEncrKeyLen,
							Value: []byte{0x10},
						},
						{
							Type:  PayloadSPPolicyParamTypeAuthAlg,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSessionAuthKeyLen,
							Value: []byte{0x14},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTCPEncrOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeSRTPAuthOffOn,
							Value: []byte{1},
						},
						{
							Type:  PayloadSPPolicyParamTypeAuthTagLen,
							Value: []byte{0x0a},
						},
					},
				},
				&PayloadKEMAC{
					SubPayloads: []*SubPayloadKeyData{
						{
							Type: 2,
							KeyData: []byte{
								0x5d, 0xb1, 0x8d, 0x95, 0x6f, 0x69, 0x67, 0xcc, 0x0d, 0x73, 0xf8, 0xb4, 0xe7, 0x76, 0xa4, 0x8a,
								0x08, 0xc4, 0x22, 0x0e, 0x52, 0x1d, 0xca, 0xa3, 0xb8, 0x00, 0xf5, 0x84, 0xfc, 0x25,
							},
						},
					},
				},
			},
		},
	},
}

func TestUnmarshal(t *testing.T) {
	for _, ca := range cases {
		t.Run(ca.name, func(t *testing.T) {
			var dec Message
			err := dec.Unmarshal(ca.enc)
			require.NoError(t, err)
			require.Equal(t, ca.dec, dec)
		})
	}
}

func TestMarshal(t *testing.T) {
	for _, ca := range cases {
		t.Run(ca.name, func(t *testing.T) {
			enc, err := ca.dec.Marshal()
			require.NoError(t, err)

			// Some test cases have 1 byte of trailing padding (0x00) that real-world
			// implementations add. We don't preserve this padding in Marshal since
			// it's not part of the spec, so trim it from expected value for comparison.
			expected := ca.enc
			if len(expected) > 0 && len(enc) == len(expected)-1 && expected[len(expected)-1] == 0x00 {
				expected = expected[:len(expected)-1]
			}

			require.Equal(t, expected, enc)
		})
	}
}

func FuzzUnmarshal(f *testing.F) {
	for _, ca := range cases {
		f.Add(ca.enc)
	}

	f.Fuzz(func(t *testing.T, b []byte) {
		var msg Message
		err := msg.Unmarshal(b)
		if err != nil {
			return
		}

		_, err = msg.Marshal()
		require.NoError(t, err)
	})
}
